# Copyright (c) 2020 Presto Labs Pte. Ltd.
# Author: taekwon

import logging
import os
import pickle

from absl import app, flags
from concurrent.futures import ProcessPoolExecutor
from coin.strategy.mm.tool.order_plot import OrderPlot

import coin.strategy.mm.tool.archive_base as abase

import pandas

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt


class BalancePlotJob(abase.SerializedJob):
  def __init__(self, baseflags, fsflags, osflags, products):
    self.baseflags = baseflags
    self.fsflags = fsflags
    self.osflags = osflags
    self.product = products[0]
    self.symbol = self.product.symbol

    if not os.path.exists("balance_plot"):
      os.makedirs("balance_plot")
    abase.SerializedJob.__init__(
        self,
        os.path.join(
            "balance_plot",
            "%s.pkl"
            % "_".join([baseflags.trading_date, self.osflags.orderlog_machine, self.symbol])))

  def run(self):
    if not os.path.exists(self.filename) or flags.FLAGS.force_rerun:
      self.serialize(self.filename)
    strat = self.deserialize()

  def serialize(self, filename):
    from coin.exchange.kr_rest.product.product_impl import generate_product_from_str2
    tds = abase.get_trading_dates(self.baseflags.trading_date)
    assert len(tds) == 1
    trading_date = tds[0]
    self.baseflags.trading_date = trading_date.strftime("%Y%m%d")

    # if self.baseflags.trading_date == "20200220":
    #   import pdb; pdb.set_trace()
    #   print('dd')

    print(self.baseflags, self.osflags)
    start_time, end_time = abase.get_time_range(self.baseflags)

    strat = OrderPlot(self.baseflags, self.osflags, self.product)
    strat.set_time_range(start_time, end_time)
    abase.run_from_order_archive_base(self.baseflags, self.osflags, on_log=strat.on_log)

    with open(filename, 'wb') as writefile:
      pickle.dump(strat, writefile)

  def deserialize_impl(self, filename):
    return pickle.load(open(filename, 'rb'))


def launch(trading_date, baseflags, fsflags, osflags, subsymbols):
  baseflags.exchange, baseflags.market_type = subsymbols[0].split(":")[:2]
  baseflags.trading_date = trading_date.strftime("%Y%m%d")
  products = abase.subsymbols_to_products(subsymbols, baseflags)
  balancejob = BalancePlotJob(baseflags, fsflags, osflags, products)
  return balancejob.run()


def run_balance_dump():
  baseflags = abase.get_base_flags()
  fsflags = abase.get_feed_subscription_flags()
  osflags = abase.get_order_subscription_flags()

  trading_dates = abase.get_trading_dates(baseflags.trading_date)

  args = tuple([baseflags, fsflags, osflags, flags.FLAGS.subsymbols.split(",")])
  print(trading_dates)
  if len(trading_dates) == 1:
    # import pdb; pdb.set_trace()
    ret_tuples = [launch(trading_dates[0], *args)]
  else:
    with ProcessPoolExecutor(max_workers=8) as executor:
      futures = []
      for trading_date in trading_dates:
        futures.append(executor.submit(launch, trading_date, *args))
      # ret_tuples = [future.result() for future in futures]


def run_plot_balance_compare(currency):
  balance1_lst = []  # coin1
  balance2_lst = []  # coin2

  # We must run BalancePlotJob Before running this.
  baseflags = abase.get_base_flags()
  trading_dates = abase.get_trading_dates(baseflags.trading_date)
  for date in trading_dates:
    fmt_date_str = date.strftime("%Y%m%d")
    logging.info("trying date " + fmt_date_str)
    # TODO(taekwon): replace this constant with flags
    try:
      sung = pickle.load(
          open(
              'balance_plot/%s_strategy-16.ap-northeast-1.huobi_BTC-USD.20200327.pkl'
              % fmt_date_str,
              'rb'))
      balance1_lst += sung._order_reader._balances[currency]
    except Exception:
      logging.warning("skip coin1 on this date")

    try:
      sung = pickle.load(
          open(
              'balance_plot/%s_strategy-41.ap-northeast-1.huobi_BTC-USD.20200327.pkl'
              % fmt_date_str,
              'rb'))
      balance2_lst += sung._order_reader._balances[currency]
    except Exception:
      logging.warning("skip coin2 on this date")

  balance1 = pandas.DataFrame(balance1_lst)
  balance2 = pandas.DataFrame(balance2_lst)

  plt.figure(figsize=(20, 10))
  plt.plot(pandas.DatetimeIndex(balance1[0]), balance1[1], 'r:', lw=0.5)
  plt.title(flags.FLAGS.subsymbols)
  plt.legend(['coin1'], loc='upper left')
  plt.twinx()
  plt.plot(pandas.DatetimeIndex(balance2[0]), balance2[1], 'b:', lw=0.5)
  plt.legend(['coin2'], loc='upper right')
  plt.savefig(os.path.join("balance_plot", f"{flags.FLAGS.subsymbols}.png"))
  plt.close()


def main(_):
  mode = flags.FLAGS.run_mode

  if mode == 'dump':
    run_balance_dump()
  elif mode == 'plot':
    currency = flags.FLAGS.subsymbols.split(":")[2].split("-")[0]
    logging.info("Plot Asset " + currency)
    run_plot_balance_compare(currency)


def init_flags():
  abase.define_base_flags()
  abase.define_feed_archive_flags()
  abase.define_order_archive_flags()

  flags.DEFINE_string('run_mode', None, 'run_mode')
  flags.DEFINE_string('subsymbols', None, 'subscript symbol name')
  flags.DEFINE_bool('force_rerun', False, '')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')


if __name__ == '__main__':
  init_flags()
  app.run(main)

# coin1
# RUST_FEED_HANDLER=1 RUST_TOPIC_RECORD=1 ./pyrunner experimental/taekwon/balance_plot.py --trading_date=20200219-20200224 --time_range=0-24 --feed_machine=feed-best.ap-northeast-1.aws --orderlog_machine=strategy-16.ap-northeast-1.huobi --subsymbols=Huobi:Futures:BTC-USD.CURRENT_QUARTER --owner=prestoinvt07 --use_feed_cache=True --run_mode=dump

# coin2
# COIN2_ORDER_LOG=1 RUST_FEED_HANDLER=1 RUST_TOPIC_RECORD=1 ./pyrunner experimental/taekwon/balance_plot.py --trading_date=20200219-20200225 --time_range=0-24 --feed_machine=feed-best.ap-northeast-1.aws --orderlog_machine=strategy-41.ap-northeast-1.huobi --subsymbols=Huobi:Futures:BTC-USD.CURRENT_QUARTER --force_rerun=True --owner=prestoinvt67 --use_feed_cache=True --run_mode=dump

# plot
# ./pyrunner experimental/taekwon/balance_plot.py --trading_date=20200219-20200225 --run_mode=plot --subsymbols=Huobi:Futures:BTC-USD.CURRENT_QUARTER
